import os
import sys
import torch
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

# 添加项目根目录到路径
sys.path.append(os.path.abspath(os.path.dirname(os.path.dirname(__file__))))

from utils.data_utils import get_data_loaders, lab_to_rgb

def show_first_image():
    """显示数据加载器加载的第一个图像对"""
    # 获取数据加载器
    data_dir = 'small_data'
    train_loader, _ = get_data_loaders(
        data_dir,
        batch_size=1,  # 每次只加载一个样本
        img_size=256,
        num_workers=0  # 不使用多线程加载
    )
    
    # 获取第一个批次
    first_batch = next(iter(train_loader))
    
    # 提取L通道和ab通道
    L = first_batch['L']
    ab = first_batch['ab']
    path = first_batch['path'][0]  # 文件路径
    
    # 打印文件路径
    print(f"加载的文件路径: {path}")
    
    # 显示灰度图像
    plt.figure(figsize=(12, 4))
    
    # 灰度图像（L通道）
    plt.subplot(1, 3, 1)
    plt.imshow(L[0, 0].numpy(), cmap='gray')
    plt.title("灰度图像 (L通道)")
    plt.axis('off')
    
    # 转换为RGB并显示
    rgb = lab_to_rgb(L, ab)[0]
    
    # 彩色图像
    plt.subplot(1, 3, 2)
    plt.imshow(rgb)
    plt.title("彩色图像 (从LAB转换)")
    plt.axis('off')
    
    # 显示原始图像
    original_img = Image.open(path).convert('RGB')
    original_img = np.array(original_img)
    
    plt.subplot(1, 3, 3)
    plt.imshow(original_img)
    plt.title("原始彩色图像")
    plt.axis('off')
    
    plt.tight_layout()
    plt.savefig('first_image.png')
    
    # 还保存一下对应的灰度图像路径
    gray_path = path.replace('.jpg', '_gray.jpg')
    print(f"对应的灰度图像路径: {gray_path}")
    
    # 检查这两个文件是否存在
    print(f"彩色图像存在: {os.path.exists(path)}")
    print(f"灰度图像存在: {os.path.exists(gray_path)}")

if __name__ == '__main__':
    show_first_image() 